#!/usr/bin/python3
# by Sun Smallwhite <niasw@pku.edu.cn>(https://github.com/niasw)

import sw.calc.calcMat
import sw.calc.gcfClus
import sw.calc.pckClus
import sw.io.loadCSV
import sw.io.saveJSON
import sw.io.adaptor
import scipy.sparse
import numpy
import json

def clusMethod(adjmat,clusnum=40,history=None,monitor=False):
  '''
assign history=[] to enable history
  '''
  eigvs=sw.calc.calcMat.maxEigVec(adjmat);
  eigvec=eigvs['eigenvector'];
  clusters=sw.calc.gcfClus.calcClusters(eigvec,clusnum=clusnum,monitor=monitor);
  if (history is None):
    clusters=sw.calc.calcMat.simplify(None,clusters);
    return clusters;
  else:
    clusters,histinit,trash=sw.calc.calcMat.simplify(None,clusters,[]);
    return clusters,histinit;

def pip_clusMethod(adjmat,numvec,lastcluslist=None,lastflowmat=None,lastnumvec=None,clusnum=40,lasthistory=None,monitor=False):
  '''
assign lasthistory=[] to enable history in initialization
adjmat is the origin flow matrix
numvec is the origin number vector 
  '''
  if ((lastflowmat is None) and (lastnumvec is None)): # need to calculate flowmat from last
    if (lastcluslist!=None): # in the middle of a pipeline
      lastflowmat,lastnumvec=sw.calc.calcMat.pip_flowMat(adjmat,numvec,lastcluslist);
    else: # first time (pipeline start)
      lastflowmat=adjmat;
      lastnumvec=numvec;
  elif (lastflowmat.shape[0]!=lastflowmat.shape[1]):
    raise(Exception("Error: LastFlowMat is not square. Row: "+lastflowmat.shape[0]+", Col: "+lastflowmat.shape[1]));
  elif (lastflowmat.shape[0]!=len(lastnumvec)):
    raise(Exception("Error: LastFlowMat and LastNumVec do not match. Mat Size: "+lastflowmat.shape[0]+", Vec Size: "+len(lastnumvec)));
  lastflowmat=scipy.sparse.coo_matrix(lastflowmat); # only sparse matrix has partly eigs solver.
  eigvs=sw.calc.calcMat.maxEigVec(lastflowmat);
  eigvec=eigvs['eigenvector'];
  clusters=sw.calc.pckClus.calcClusters(eigvec,lastnumvec,clusnum=clusnum,monitor=monitor);
  if (lasthistory is None): # history flag off
    if (lastcluslist!=None): # in the middle of a pipeline
      newcluslist=sw.calc.calcMat.simplify(lastcluslist,clusters);
    else: # first time
      newcluslist=clusters;
    return newcluslist;
  else: # history flag on
    newcluslist,newhistory,lasthistory=sw.calc.calcMat.simplify(lastcluslist,clusters,lasthistory);
    return newcluslist,newhistory,lasthistory;

def Adaptor4VEGAS(in_linkfile,in_numvecfile,out_clusprefix,clusnumlist,indexstart=0,withdata=False,history=False,monitor=False,localoutput=False):
  '''
# adaptor to adapt Xuan Yuan's IO from VEGAS
#
# input:
#  in_linkfile: the file stores links. (sparse matrix coo format)
#  in_numvecfile: the file stores number vector. (column csv)
#  out_clusprefix: output filename prefix
#  clusnumlist: the list of cluster number requests
#  indexstart: index start from 1
#  withdata: in_linkfile contains value column
#  history: whether generate history json data
#  monitor: see eigenvector figures and error reports
#  localoutput: create data file for Sun Sibai's preview engine
# output:
#  write clustering result into files: out_clusprefix.clustering_times.paper_number.cluster_number
  '''
  print('Loading Link Data & Number Vector Data ...');
  adjmat=sw.io.loadCSV.linkCSV2adjMat('./dat/'+in_linkfile,indexstart=indexstart,withdata=withdata);
  print('Total Node Number: '+str(adjmat.shape[0]));
  if (withdata):
    numvec=sw.io.loadCSV.loadCSVvector('./dat/'+in_numvecfile);
  else:
    numvec=numpy.array([1 for it in range(0,adjmat.shape[0])]);
  print('Total Paper Number: '+str(numvec.sum()));
  if (adjmat.shape[0]!=adjmat.shape[1]):
    raise(Exception('Adjacent Matrix is not square. Row:'+str(adjmat.shape[0])+' ,Col:'+str(adjmat.shape[1])));
  if (adjmat.shape[0]!=len(numvec)):
    raise(Exception('Dimension of number vector and adjacent matrix do not match. Mat:'+str(adjmat.shape[0])+' ,Vec:'+str(len(numvec))));
  cluslist=None; # first time: cluslist=None
  histlist=[];
  flowmat=None;
  curnumvec=None;
  for it in range(0,len(clusnumlist)):
   clusnum=clusnumlist[it];
   print('Stair Clustering ['+str(clusnum)+','+str(it+1)+'/'+str(len(clusnumlist))+'] ...');
   if (history): # history flag on
     if (cluslist is None): # first time
      if (withdata): # origin data with numvec, use incremental algorithm only
       cluslist,histnew,histold=pip_clusMethod(adjmat,numvec,cluslist,lastflowmat=flowmat,lastnumvec=curnumvec,clusnum=clusnum,lasthistory=[],monitor=monitor);
       histlist=[histnew];
      else: # origin data without numvec, use normal algorithm to speed up
       cluslist,histinit=clusMethod(adjmat,clusnum=clusnum,history=[],monitor=monitor);
       histlist=[histinit];
     else: # in the middle of a pipeline
       cluslist,histnew,histold=pip_clusMethod(adjmat,numvec,cluslist,lastflowmat=flowmat,lastnumvec=curnumvec,clusnum=clusnum,lasthistory=histlist[0],monitor=monitor);
       histlist[0]=histold;
       histlist.insert(0,histnew);
   else: # history flag off
     if (cluslist is None): # first time
       cluslist=clusMethod(adjmat,clusnum=clusnum,monitor=monitor);
     else: # in the middle of a pipeline
       cluslist=pip_clusMethod(adjmat,numvec,cluslist,lastflowmat=flowmat,lastnumvec=curnumvec,clusnum=clusnum,monitor=monitor);
   if (clusnum<len(numvec)): # otherwise the clusters do not change.
     flowmat,curnumvec=sw.calc.calcMat.pip_flowMat(adjmat,numvec,cluslist);
     if monitor:
       print('-> Number of Clusters = '+str(len(cluslist)));
       print('-> Total Flow = '+str(flowmat.sum()));
       print('-> Total Flow Square = '+str(numpy.linalg.norm(flowmat,ord='fro')**2));
     if (localoutput):
       sw.io.saveJSON.saveJSON('./out/'+out_clusprefix+'_clus.'+str(it+1)+'.'+str(numvec.sum())+'.'+str(clusnum)+'.json',cluslist);
       sw.io.saveJSON.saveJSON('./out/'+out_clusprefix+'_flow.'+str(it+1)+'.'+str(numvec.sum())+'.'+str(clusnum)+'.json',flowmat.tolist());
   #clusvec=sw.io.adaptor.clus2vec(cluslist,indexstart);
   #sw.io.adaptor.vec2csv(clusvec,out_clusprefix+'.'+str(it+1)+'.'+str(numvec.sum())+'.'+str(clusnum));
  if (history):
   histlist=sw.io.adaptor.histIdxMap(histlist,indexstart);
   sw.io.saveJSON.saveJSON('./out/'+out_clusprefix+'_hist.'+str(numvec.sum())+'.'+str(clusnumlist[0])+'.json',histlist);

if (__name__=='__main__'):
  '''
# python adaptVEGAS.py in_linkfile in_numvecfile out_clusprefix clusnumlist indexstart=0 withdata=True history=False monitor=False localoutput=False
  '''
  import sys
  args=[None,'','','',[],0,'True','False','False','False'];
  for it in range(0,len(sys.argv)):
    args[it]=sys.argv[it];
  Adaptor4VEGAS(args[1],args[2],args[3],json.loads(args[4]),int(args[5]),args[6]=='True',args[7]=='True',args[8]=='True',args[9]=='True');
